import requests
import utils
from Crypto.Cipher import PKCS1_OAEP
from Crypto.PublicKey import RSA
from Crypto.Hash import SHA256
import binascii
import http.cookies

headers = {
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko)'
                  ' Chrome/114.0.0.0 Safari/537.36'
}


class BiliCookieRefresher:
    def __init__(self, initial_cookie, initial_refresh_token):
        """
        初始化B站Cookie刷新器

        :param initial_cookie: 初始Cookie字典
        :param initial_refresh_token: 初始refresh_token
        """
        self.simple_cookie = http.cookies.SimpleCookie()
        self.cookie = initial_cookie
        self.refresh_token = initial_refresh_token
        self.session = requests.Session()
        self.session.cookies.update(initial_cookie)

        # RSA公钥用于生成CorrespondPath
        self.public_key = RSA.importKey('''\
-----BEGIN PUBLIC KEY-----
MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQDLgd2OAkcGVtoE3ThUREbio0Eg
Uc/prcajMKXvkCKFCWhJYJcLkcM2DKKcSeFpD/j6Boy538YXnR6VhcuUJOhH2x71
nzPjfdTcqMz7djHum0qSZA0AyCBDABUqCrfNgCiJ00Ra7GmRj+YCK1NJEuewlb40
JNrRuoEUXpabUzGB8QIDAQAB
-----END PUBLIC KEY-----''')

    def check_refresh_needed(self):
        """
        检查是否需要刷新Cookie

        :return: (是否需要刷新, 时间戳) 元组
        """
        url = "https://passport.bilibili.com/x/passport-login/web/cookie/info"
        params = {
            'csrf': self.cookie.get('bili_jct', '')
        }

        try:
            response = self.session.get(url, headers=headers, params=params)
            data = response.json()
            if data['code'] == 0:
                return data['data']['refresh'], data['data']['timestamp']
            else:
                print(f"检查刷新失败: {data['message']}")
                return False, None
        except Exception as e:
            print(f"检查刷新异常: {str(e)}")
            return False, None

    def generate_correspond_path(self, timestamp):
        """
        生成CorrespondPath

        :param timestamp: 当前毫秒时间戳
        :return: correspond_path字符串
        """
        try:
            cipher = PKCS1_OAEP.new(self.public_key, SHA256)
            encrypted = cipher.encrypt(f'refresh_{timestamp}'.encode())
            return binascii.b2a_hex(encrypted).decode()
        except Exception as e:
            print(f"生成CorrespondPath异常: {str(e)}")
            return None

    def get_refresh_csrf(self, correspond_path):
        """
        获取refresh_csrf

        :param correspond_path: 生成的correspond_path
        :return: refresh_csrf字符串
        """
        url = f"https://www.bilibili.com/correspond/1/{correspond_path}"

        try:
            response = self.session.get(url, headers={
                **utils.request.BILIBILI_COMMON_HEADERS
            })
            if response.status_code == 200:
                # 从HTML中提取refresh_csrf
                # 这里简化处理，实际应该用HTML解析器如BeautifulSoup
                start = response.text.find('id="1-name">') + len('id="1-name">')
                end = response.text.find('</div>', start)
                refresh_csrf = response.text[start:end].strip()
                return refresh_csrf
            else:
                print(f"获取refresh_csrf失败，HTTP状态码: {response.status_code}")
                return None
        except Exception as e:
            print(f"获取refresh_csrf异常: {str(e)}")
            return None

    def refresh_cookie(self):
        """
        执行完整的Cookie刷新流程

        :return: 是否成功刷新
        """
        # 1. 检查是否需要刷新
        need_refresh, timestamp = self.check_refresh_needed()
        if not need_refresh:
            print("当前不需要刷新Cookie")
            return True

        print("检测到需要刷新Cookie，开始刷新流程...")

        # 2. 生成CorrespondPath
        correspond_path = self.generate_correspond_path(timestamp)
        if not correspond_path:
            print("生成CorrespondPath失败")
            return False

        # 3. 获取refresh_csrf
        refresh_csrf = self.get_refresh_csrf(correspond_path)
        if not refresh_csrf:
            print("获取refresh_csrf失败")
            return False

        # 保存旧的refresh_token用于后续确认
        old_refresh_token = self.refresh_token

        # 4. 刷新Cookie
        url = "https://passport.bilibili.com/x/passport-login/web/cookie/refresh"
        data = {
            'csrf': self.cookie.get('bili_jct', ''),
            'refresh_csrf': refresh_csrf,
            'source': 'main_web',
            'refresh_token': self.refresh_token
        }

        try:
            response = self.session.post(url, headers=headers, data=data)
            result = response.json()
            if result['code'] == 0:
                # 更新refresh_token
                self.refresh_token = result['data']['refresh_token']
                self.simple_cookie = response.cookies
                # 更新session的cookies
                new_cookies = response.cookies.get_dict()
                self.cookie.update(new_cookies)
                self.session.cookies.update(new_cookies)
                print("Cookie刷新成功")

                # 5. 确认更新
                return self.confirm_refresh(old_refresh_token)
            else:
                print(f"刷新Cookie失败: {result['message']}")
                return False
        except Exception as e:
            print(f"刷新Cookie异常: {str(e)}")
            return False

    def confirm_refresh(self, old_refresh_token):
        """
        确认更新，使旧的refresh_token失效

        :param old_refresh_token: 旧的refresh_token
        :return: 是否成功确认
        """
        url = "https://passport.bilibili.com/x/passport-login/web/confirm/refresh"
        data = {
            'csrf': self.cookie.get('bili_jct', ''),
            'refresh_token': old_refresh_token
        }

        try:
            response = self.session.post(url, headers=headers, data=data)
            result = response.json()
            if result['code'] == 0:
                print("确认更新成功")
                return True
            else:
                print(f"确认更新失败: {result['message']}")
                return False
        except Exception as e:
            print(f"确认更新异常: {str(e)}")
            return False

    def get_current_cookie(self):
        """
        获取当前Cookie

        :return: 当前Cookie字典
        """
        return self.cookie

    def get_current_refresh_token(self):
        """
        获取当前refresh_token

        :return: 当前refresh_token
        """
        return self.refresh_token

    def get_simple_cookie(self):
        """
        获取当前simple_cookie

        :return: 当前simple_cookie
        """
        return self.simple_cookie


# 使用示例
if __name__ == "__main__":
    # 初始化时需要提供初始Cookie和refresh_token
    # 这些值应该从登录后的响应中获取并保存
    initial_cookie = {
        'SESSDATA': 'd927476c%2C1769583733%2Cc7681%2A81CjA0VGoJQxNXZSuvVvRP7RWLd3T47HnzrvhqagL3UqcqlaFQwTYWpI5jAkhPuB4W-GYSVk1xcXZ3VHlRbzRkTGpTMHAxVC1lWi1UZThLbXpESmFPQlp1QURKUW9pY3o3ajc3MXFpSDZqY0tNZm0tQjFmX1NVZWlKZHFTOVB2Yms1QU5QWmFwZUlBIIEC',
        'bili_jct': '8d4cbb0e635903c0b35c58dc00fa3960',
        'DedeUserID': '107817747',
        'DedeUserID__ckMd5': '40c209af636e9508',
        'sid': 'hjla8ebp'
    }

    initial_refresh_token = '98bce25a687a6aad4aca1c06b82dc781'

    # 创建刷新器实例
    refresher = BiliCookieRefresher(initial_cookie, initial_refresh_token)

    # 执行刷新流程
    success = refresher.refresh_cookie()

    if success:
        print("Cookie刷新流程完成")
        print("新的Cookie:", refresher.get_current_cookie())
        print("新的refresh_token:", refresher.get_current_refresh_token())
    else:
        print("Cookie刷新流程失败")
